from typing import Callable

import torch
from torch import nn


class CBA(nn.Module):
    def __init__(
        self,
        ch_in,
        ch_out,
        kernel=3,
        stride=1,
        groups=1,
        activation: Callable[..., nn.Module] = nn.ReLU,
    ):
        super(CBA, self).__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(
                ch_in, ch_out, kernel, stride, padding=kernel // 2, groups=groups, bias=False
            ),
            nn.BatchNorm2d(ch_out),
            activation(),
        )

    def forward(self, x):
        return self.layers(x)


class ReChannel(nn.Module):
    def __init__(self, inc, ouc, act_op):
        super().__init__()
        self.module_list = nn.ModuleList(
            [CBA(inc[i], ouc[i], 1, activation=act_op) for i in range(len(inc))]
        )

    def forward(self, x):
        return [module(p) for module, p in zip(self.module_list, x)]


class UNode(nn.Module):
    def __init__(self, inc1, inc2, outc, fusion, act_op):
        super().__init__()
        if fusion == "add":
            self.cba = CBA(inc1, outc, 3, activation=act_op)
        elif fusion == "cat":
            self.cba = CBA(inc1 + inc2, outc, 3, activation=act_op)
        self.act = act_op(inplace=True)
        self.up_sample = nn.Upsample(scale_factor=2)
        self.fusion = fusion

    def forward(self, xi, xu):
        if self.fusion == "add":
            x = xi + self.up_sample(xu)
        elif self.fusion == "cat":
            x = torch.cat([xi, self.up_sample(xu)], 1)
        return self.cba(self.act(x))


class UBlock(nn.Module):
    def __init__(self, incs, act_op, fusion):
        super().__init__()
        self.nodes = nn.ModuleList(
            [
                UNode(incs[i], incs[i - 1], incs[i - 1], fusion, act_op)
                for i in range(len(incs) - 1, 0, -1)
            ]
        )

    def forward(self, x):
        outs = [x[-1]]
        for node, x_i in zip(self.nodes, x[-2::-1]):
            outs.append(node(x_i, outs[-1]))
        return outs
